/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;

import alloy.ast.ASTDepthFirstReplacer;
import alloy.ast.CardinalityExpr;
import alloy.ast.CompOp;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.Decls;
import alloy.ast.ElemFormula;
import alloy.ast.Expr;
import alloy.ast.Formula;
import alloy.ast.Id;
import alloy.ast.LeafId;
import alloy.ast.MultiplicityExpr;
import alloy.ast.Qualifier;
import alloy.ast.Qualifiers;
import alloy.ast.QuantifiedExpr;
import alloy.ast.QuantifiedFormula;
import alloy.ast.Quantifier;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.ast.Variables;
import alloy.semantic.LocalScope;
import alloy.type.UnionType;
import alloy.util.Dbg;

/**
 * <p>
 * A visitor which desugars quantifiers. The quantifiers no, lone, and one
 * desugar as follows:
 * <ul>
 * <li>no [disj] x1, x2, ..., xn : m S | F(x1, ..., xn) ==> all [disj] x1, x2,
 * ..., xn : m S | !F(x1, ..., xn)
 * <li>lone [disj] x1, x2, ..., xn : m S | F(x1, ..., xn) ==> all [disj] x1,
 * x2, ..., xn : m S | F(x1, ..., xn) => (all [disj] y1, y2, ...., yn : m S |
 * F(y1, ..., yn) => (x1 = y1) && ... && (xn = yn)"
 * <li>one [disj] x1, x2, ..., xn : m S | F(x1, ..., xn) ==> some [disj] x1,
 * x2, ..., xn : m S | F(x1, ..., xn) && (all [disj] y1, y2, ...., yn : m S |
 * F(y1, ..., yn) => (x1 = y1) && ... && (xn = yn)
 * </ul>
 * <ul>
 * </ul>
 * <p>
 * Quantified expressions are as follows:
 * <ul>
 * <li>no expr ==> # expr = 0
 * <li>lone expr ==> # expr <= 1
 * <li>one expr ==> # expr = 1
 * <li>some expr ==> # expr > 0
 * </ul>
 */

public class DesugarQuantifiersVisitor extends ASTDepthFirstReplacer {
    
    public Object visit(QuantifiedExpr quantExpr_) {
		// first, apply visitor to expr
		quantExpr_.setExpr((Expr) quantExpr_.getExpr().applyReturnVisitor(this));
		Quantifier quant = quantExpr_.getQuantifier();
		Formula replacement = _desugarQuantifier(quant, (Expr) quantExpr_.getExpr());
		replacement.annotateTransformReplacing("quantifier " + quant.nodeString() + " desugared", quantExpr_);
		return replacement;
	}
    
    public Object visit(QuantifiedFormula quantFormula_) {
        // first, desugar decls and body
        quantFormula_.setDecls((Decls)quantFormula_.getDecls().applyReturnVisitor(this));
        quantFormula_.setFormula((Formula)quantFormula_.getFormula().applyReturnVisitor(this));
        Formula replacement = null;
        
        // don't desugar all or some for formulas
        Quantifier quant = quantFormula_.getQuantifier();
        if (quant.equals(Quantifier.ALL) || quant.equals(Quantifier.SOME)) {
            replacement = quantFormula_;
        } else {
            Formula body = quantFormula_.getFormula();
            Decls decls = quantFormula_.getDecls();
            
            if (quant.equals(Quantifier.NO)) {
                replacement = new QuantifiedFormula(Quantifier.ALL,
                        decls,
                        body.not(), true);
            } else {
                if (higherOrder(decls)) { // general desugaring if we have higher
                    // order decls
                    if (quant.equals(Quantifier.LONE)) {
                        replacement = new QuantifiedFormula(Quantifier.ALL,
                                decls,
                                body.implies(_genNestedFormula(quantFormula_)), true); 
                    } else { // opcode is Quantifier.ONE
                        replacement = new QuantifiedFormula(Quantifier.SOME,
                                decls,
                                body.and(_genNestedFormula(quantFormula_)), true);
                    }
                } else { // optimize for the common case when we don't
                    replacement = desugarAsComprehension(quantFormula_);
                }
            }
        }

            return replacement ; // _breakUpQuantifiers(replacement);
        }    
        
        // returns true if there is at least one decl in decls that
        // is higher order, otherwise returns false
        private static boolean higherOrder(Decls decls) {
            for (Iterator iter = decls.getDeclIter(); iter.hasNext();) {
                if (((Decl) iter.next()).isHigherOrder())
                    return true;
            }
            return false;
        }
        
        // desugars the specified quantified formula (with quantifier being one or
        // lone)
        // as a set comprehension:
        // lone x | F(x) ==> # { x | F(x) } <= 1
        // one x | F(x) ==> # { x | F(x) } = 1
        private static Formula desugarAsComprehension(
                QuantifiedFormula quantFormula_) {
            // create comprehension expression
            Decls decls = quantFormula_.getDecls();
            ComprehensionExpr compExpr = new ComprehensionExpr(decls, quantFormula_
                    .getFormula(), true);
            // set type of new comprehension expression by iterating through decls
            // and creating new union type
            LocalScope quantFormScope = quantFormula_.getLocalScope();
            Iterator declIter = decls.getDeclIter();
            UnionType uType = null;
            while (declIter.hasNext()) {
                Decl decl = (Decl) declIter.next();
                Iterator varIter = decl.getVariables().getVariableIter();
                while (varIter.hasNext()) {
                    Variable var = (Variable) varIter.next();
                    UnionType curType = quantFormScope.lookupType(var.nodeString());
                    uType = (uType == null ? curType : uType.product(curType));
                }
            }
            compExpr.setType(uType);
            Formula replacement = _desugarQuantifier(quantFormula_.getQuantifier(),
                    (Expr) compExpr);
            replacement.annotateTransformReplacing("quantifier \'"
                    + quantFormula_.getQuantifier().nodeString() + "\' desugared",
                    quantFormula_);
            return replacement;
            
        }

        /**
	 	* For a formula
	 	*   "q [disj] x1, x2, ..., xn : m S | F(x1, ..., xn)"
	 	* generate a formula
	 	*   "(all [disj] y1, y2, ...., yn : m S | F(y1, ..., yn) => (x1 = y1) && ... && (xn = yn)"
	 	* 
	 	* @param quantFormula_
	 	* @return
	 	*/
	private Formula _genNestedFormula(QuantifiedFormula quantFormula_) {
		Decls originalDecls = quantFormula_.getDecls();
		Decls genDecls = new Decls(originalDecls.getLocation());
		Map originalToGen = new HashMap();

		// generate the declaration [disj] y1, y2, ..., yn : m S
		// and the formula (x1 = y1 && ... && xn = yn)
		Formula consequent = null;

		for (Iterator declIter = originalDecls.getDeclIter(); declIter.hasNext();) {
			Decl originalDecl = (Decl) declIter.next();
			Variables originalVariables = originalDecl.getVariables();

			Variables genVariables = new Variables(originalVariables.getLocation());

			for (Iterator varIter = originalVariables.getVariableIter(); varIter.hasNext();) {
				// get xi
				Variable originalVar = (Variable) varIter.next();

				// generate yi
				Variable genVar = new Variable(originalVar.getLocation(),
						Id.generateId(originalVar.getId().nodeString()), false);
				// add yi to variables y1, y2, ..., yn
				genVariables.addVariable(genVar);

				// create variable expr for xi
				VariableExpr originalVarExpr = new VariableExpr((Variable) originalVar.copy());
				originalVarExpr.setType(originalDecl.getExpr().getType());

				// create variable expr for yi
				VariableExpr genVarExpr = new VariableExpr((Variable) genVar.copy());
				genVarExpr.setType(originalVarExpr.getType());

				// add xi = yi to (x1 = y1) && ... && (xn = yn)
				Formula xEQy = new ElemFormula(originalVar.getLocation(), originalVarExpr, CompOp.EQUALS, genVarExpr);
				consequent = (consequent == null ? xEQy : consequent.and(xEQy));

				// add the mapping from xi to variables expr for yi to
				// originalToGen
				originalToGen.put(originalVar.getId().nodeString(), genVarExpr.copy());

			}

			// generate qualifiers [disj]
			Qualifiers genQualifiers = new Qualifiers();
			if (originalDecl.isDisj())
				genQualifiers.addQualifier(Qualifier.DISJ);
			if (originalDecl.isExh())
				genQualifiers.addQualifier(Qualifier.EXH);

			// generate the declaration [disj] y1, y2, ..., yn : m S
			Decl genDecl = new Decl(originalDecl.getLocation(), genQualifiers,
					genVariables, (MultiplicityExpr) originalDecl.getExpr()	.copy());

			genDecls.addDecl(genDecl);
		}

		// generate F(y1, ..., yn)
		Formula antecedent = (Formula) quantFormula_.getFormula().copy();
		antecedent.applyReturnVisitor(new ExprAndTypeReplaceVisitor(
				originalToGen, "original variable", "generated variable"));

		// generate "(all [disj] y1, y2, ...., yn : m S | F(y1, ..., yn) => (x1
		// = y1) && ... && (xn = yn)"
		return new QuantifiedFormula(quantFormula_.getLocation(),
				Quantifier.ALL, genDecls, antecedent.implies(consequent), true);

	}
           
        /**
		 * Creates a formula expressing the constraint of the quantifier on the
		 * expression.
		 * 
		 * @param quantifier_
		 *            the quantifier
		 * @param expr_
		 *            the expression to be constrained
		 * 
		 * @return the formula expressing the constraint of the quantifier on
		 *         the expression
		 */
        
        private static Formula _desugarQuantifier(Quantifier quantifier_, Expr expr_) {
            // create the cardinality expression
            CardinalityExpr cardExpr = new CardinalityExpr(expr_);
            
            if (quantifier_.equals(Quantifier.NO)) {
                Expr empty = UnionType.empty(expr_.getType().arity()).toExpr();
                return expr_.in(empty);
            } else if (quantifier_.equals(Quantifier.LONE)) {
	            return QuantifiedExpr.lone(expr_);
            } else if (quantifier_.equals(Quantifier.ONE)) {
                Formula f = expr_.in(UnionType.empty(expr_.getType().arity()).toExpr()).not();
            	return f.and(QuantifiedExpr.lone((Expr) expr_.copy()));
            } else /*if (quantifier_.equals(Quantifier.SOME))*/ {
                return expr_.in(UnionType.empty(expr_.getType().arity()).toExpr()).not();
            }
        }
        
        /**
         * break up quantifiers, so that all a : A, b : B | F(a,b) becomes all a : A |
         * all b : B | F(a,b). useful for optimizations
         */
        private Formula _breakUpQuantifiers(QuantifiedFormula quantFormula_) {
            // don't separate multiple times
            if (quantFormula_.isSeparated)
                return quantFormula_;
            
            LeafIdMap leafIdMap = quantFormula_.getLeafIdMap();
            List newDecls = new ArrayList();
            // create a flat list of decls for the quantified formulas
            Iterator declIter = quantFormula_.getDecls().getDeclIter();
            while (declIter.hasNext()) {
                Decl curDecl = (Decl) declIter.next();
                // don't do any desugaring if there are qualifiers
                if (curDecl.isDisj() || curDecl.isExh()) {
                    newDecls.add(curDecl);
                    continue;
                }
                // otherwise, add a decl for each variable
                MultiplicityExpr multExpr = curDecl.getExpr();
                Iterator varIter = curDecl.getVariables().getVariableIter();
                while (varIter.hasNext()) {
                    Variable curVar = (Variable) varIter.next();
                    Decl newDecl;
                    newDecl = Decl.in(curVar, (MultiplicityExpr) multExpr.copy());
                    newDecl.setLocation(curDecl.getLocation());
                    newDecls.add(newDecl);
                }
            }
            
            // now, create the new quantified formulas
            Formula curFormula = (Formula) quantFormula_.getFormula().copy();
            Quantifier quant = quantFormula_.getQuantifier();
            Dbg.chk(quant.equals(Quantifier.ALL) || quant.equals(Quantifier.SOME));
            ListIterator newDeclIter = newDecls.listIterator(newDecls.size());
            while (newDeclIter.hasPrevious()) {
                Formula prevFormula = curFormula;
                Decl decl = (Decl) newDeclIter.previous();
                if (quant.equals(Quantifier.ALL)) {
                    curFormula = QuantifiedFormula.all(decl, curFormula);
                } else {
                    curFormula = QuantifiedFormula.some(decl, curFormula);
                }
                ((QuantifiedFormula) curFormula).isSeparated = true;
                if (leafIdMap != null) {
                    // create a new leaf id map for the new formula
                    LeafIdMap newMap = new LeafIdMap(null);
                    Iterator varIter = decl.getVariables().getVariableIter();
                    // need to iterate since there still might be multiple variables
                    // (if we didn't break them up because of a qualifier)
                    while (varIter.hasNext()) {
                        String varString = ((Variable) varIter.next()).nodeString();
                        LeafId leafId = leafIdMap.lookupLeafId(varString);
                        Dbg.chk(leafId);
                        newMap.put(varString, leafId);
                    }
                    ((QuantifiedFormula) curFormula).setLeafIdMap(newMap);
                }
                if (prevFormula instanceof QuantifiedFormula) {
                    // set the parent pointer of the local scope to the
                    // new local scope
                    LocalScope outer = ((QuantifiedFormula) curFormula)
                    .getLocalScope();
                    LocalScope inner = ((QuantifiedFormula) prevFormula)
                    .getLocalScope();
                    inner.addParent(outer);
                    // same for the leaf id map
                    if (leafIdMap != null) {
                        LeafIdMap outerMap = ((QuantifiedFormula) curFormula)
                        .getLeafIdMap();
                        LeafIdMap innerMap = ((QuantifiedFormula) prevFormula)
                        .getLeafIdMap();
                        innerMap.setParent(outerMap);
                    }
                }
                
            }
            // finally, set the last formula's local scope's parent
            // to be the parent of the original formula, if one
            // exists
            Iterator parents = quantFormula_.getLocalScope().getParents();
            if (parents.hasNext()) {
                LocalScope parent = (LocalScope) parents.next();
                ((QuantifiedFormula) curFormula).getLocalScope().addParent(parent);
            }
            // same for leaf id map
            if (leafIdMap != null) {
                LeafIdMap parentMap = leafIdMap.getParent();
                if (parentMap != null) {
                    ((QuantifiedFormula) curFormula).getLeafIdMap().setParent(
                            parentMap);
                }
            }
            curFormula.annotateTransformReplacing(
                    "variables in quantifier separated", quantFormula_);
            return curFormula;
        }
        
    }